import glob
import os

import cv2
import mmcv
import numpy as np

source_path = '/mnt/data302/xulongwei/new_work/jpg_datasets/LiTS'
target_path = '/mnt/data302/xulongwei/new_work/jpg_datasets/LiTS-reduce_zero'
if not os.path.exists(target_path):
    os.mkdir(target_path)


def reduce_zero(source_path, target_path):
    imgs = glob.glob(os.path.join(source_path, 'train', 'Labels', '*'))
    img_full_zero = []  # 整个label都是0的
    img_least_one = []  # label中至少有一个像素点不为0
    for img in imgs:
        img_buff = open(img, 'rb').read()
        img_array = mmcv.imfrombytes(img_buff, 'unchanged')
        if img_array.max() != 0:
            img_least_one.append(img)
        else:
            img_full_zero.append(img)
    img_full_zero = np.array(img_full_zero)
    img_least_one = np.array(img_least_one)
    np.save(os.path.join(target_path, 'full_zero.npy'), img_full_zero)
    np.save(os.path.join(target_path, 'nonfull_zero.npy'), img_least_one)
    with open(os.path.join(target_path, 'process.log'), 'a+') as f:
        f.write('Original Dataset has {0} CT scans\n'
                'There are {1} CT scans without liver\n'
                'full_zero.npy recorded CT scans full of zero label saved at {2}\n'
                'nonfull_zero.npy recorded CT scans at least have liver saved at{3}\n'
                .format(len(imgs), len(img_full_zero), os.path.join(target_path, 'full_zero.npy'),
                        os.path.join(target_path, 'nonfull_zero.npy')))
    return img_least_one


def save_nonezero(target_path, img_least_one):
    for img in img_least_one:
        if not os.path.exists(os.path.join(target_path, img.split('/')[-3])):
            os.mkdir(os.path.join(target_path, img.split('/')[-3]))
        if not os.path.exists(os.path.join(target_path, img.split('/')[-3], img.split('/')[-2])):
            os.mkdir(os.path.join(target_path, img.split('/')[-3], img.split('/')[-2]))
        img_buff = open(img, 'rb').read()
        img_array = mmcv.imfrombytes(img_buff, 'unchanged')
        cv2.imwrite(os.path.join(target_path, img.split('/')[-3], img.split('/')[-2],
                                 img.split('/')[-1]), img_array)


def makedemodatasets(source_path, target_path, num):
    if not os.path.exists(target_path):
        os.mkdir(target_path)
    source_path_new = os.path.join(source_path, 'train', 'Images')
    imgs = glob.glob(os.path.join(source_path_new, '*'))
    assert num <= len(imgs)
    del_index = np.random.randint(0, len(imgs), num)
    for count, ind in enumerate(del_index):
        img = imgs[ind]
        img_buff = open(img, 'rb').read()
        image = mmcv.imfrombytes(img_buff, 'unchanged')
        if not os.path.exists(os.path.join(target_path, img.split('/')[-3])):
            os.mkdir(os.path.join(target_path, img.split('/')[-3]))
        if not os.path.exists(os.path.join(target_path, img.split('/')[-3], img.split('/')[-2])):
            os.mkdir(os.path.join(target_path, img.split('/')[-3], img.split('/')[-2]))
        cv2.imwrite(img.replace(source_path, target_path), image)
        lab = (os.path.join(img.replace('Images', 'Labels').replace('.jpg', '.png')))
        label_buff = open(lab, 'rb').read()
        label = mmcv.imfrombytes(label_buff, 'unchanged')
        if not os.path.exists(os.path.join(target_path, lab.split('/')[-3])):
            os.mkdir(os.path.join(target_path, lab.split('/')[-3]))
        if not os.path.exists(os.path.join(target_path, lab.split('/')[-3], lab.split('/')[-2])):
            os.mkdir(os.path.join(target_path, lab.split('/')[-3], lab.split('/')[-2]))
        cv2.imwrite(lab.replace(source_path, target_path), label)
        if count < 100:
            cv2.imwrite(img.replace(source_path, target_path).replace('train', 'val'), image)
            cv2.imwrite(lab.replace(source_path, target_path).replace('train', 'val'), label)


# img_least_one_label = reduce_zero(source_path, target_path)
# save_nonezero(target_path, img_least_one_label)
# img_least_one_image = [i.replace('.png', '.jpg').replace('Labels', 'Images') for i in img_least_one_label]
# save_nonezero(target_path, img_least_one_image)
makedemodatasets(target_path,
                 '/mnt/data302/xulongwei/new_work/jpg_datasets/demo-LiTS-reduce_zero', 1000)
